/*
 * COPYRIGHT:	2010 Christopher Smith, licenced under the GNU GPL v2; see COPYING in the top level directory
 * PROJECT:		sdedit
 * FILE:		ntds ldap.cpp
 * PURPOSE:		NTDS (NT Directory Services (aka Active Directory)) over LDAP (the Lightweight Directory Access Protocol) common support
 */

namespace Ldap
{

class NtdsLdapModuleBase : public Module::ModuleObject 
{
public:
	auto_SDE_OPTS ModOpts;

	NtdsLdapModuleBase(SDE_OPTS *TemplateOpts) :
		ModOpts(*TemplateOpts)
	{}
};

void RaiseLdapError(ULONG Err, BOOLEAN Noncontinuable = TRUE)
{
	PCWSTR ErrString = ldap_err2string(Err);
	RaiseError(STATUS_LDAP_ERROR, Err, L"An LDAP Error Has Occured", ErrString, Noncontinuable);
}

template <class T> T RaiseOnLdapErrorT(T Result, T ErrorResult, BOOLEAN Noncontinuable = TRUE)
{
	if(Result == ErrorResult)
		RaiseLdapError(LdapGetLastError(), Noncontinuable);
	return Result;
}

ULONG RaiseOnLdapError(ULONG Result, BOOLEAN Noncontinuable = TRUE)
{
	if(Result != LDAP_SUCCESS)
		RaiseLdapError(Result, Noncontinuable);
	return Result;
}

void free_ldap_message(LDAPMessage *msg)
{
	ldap_msgfree(msg);
}

void free_ldap_value_len(berval **msg)
{
	ldap_value_free_len(msg);
}

void free_ldap_value(WCHAR **msg)
{
	ldap_value_free(msg);
}

void free_ldap_connection(LDAP *ld)
{
	ldap_unbind(ld);
}

void free_BERVAL(berval *bv)
{
	ber_bvfree(bv);
}

void free_BerElement1(BerElement *bv)
{
	ber_free(bv, 1);
}

//a noncritical hint to the server that the query will be satisfied without need of referrals (to other DCs or partitions)
//used to avoid timeouts associated with referrals to unreachable domains when referrals shouldn't be necessary in the first place
LDAPControl DomainScopeHintCtrl = 
{
	LDAP_SERVER_DOMAIN_SCOPE_OID_W,
	{0, NULL},
	FALSE
};

auto_take<WCHAR> DomainDnsToDc(PCWSTR Dns)
{
	if(!Dns)
		return auto_take<WCHAR>(NULL);

	SIZE_T DnsLen = wcslen(Dns);
	auto_szpod<WCHAR> Ret = auto_szpod_empty_ct(DnsLen * 3 + 3); //worst-case
	PCWSTR DnsPtStart = Dns;
	PCWSTR DnsMax = Dns + DnsLen;
	PWSTR OutPos = Ret;
	BOOLEAN FirstComponent = TRUE;

	for(PCWSTR i = Dns; i <= DnsMax; i++)
	{
		if(*i == '.' || *i == 0)
		{
			if(!FirstComponent)
				*(OutPos++) = L',';
			*(OutPos++) = L'D';
			*(OutPos++) = L'C';
			*(OutPos++) = L'=';
			SIZE_T PtLen = reinterpret_cast<ULONG_PTR>(i) - reinterpret_cast<ULONG_PTR>(DnsPtStart);
			RtlCopyMemory(OutPos, DnsPtStart, PtLen);
			OutPos += PtLen / sizeof(WCHAR);
			DnsPtStart = ++i;
			FirstComponent = FALSE;
		}
	}

	return Ret.release();
}

auto_take<WCHAR> DnToDomainDns(PCWSTR Dn)
{
	auto_free<PWSTR, free_ldap_value> Parts = auto_take_(ldap_explode_dn(const_cast<WCHAR *const>(Dn), FALSE));

	PWSTR *Part = Parts;

	//find first domain component
	do
		if(((*Part)[0] == L'D' || (*Part)[0] == L'd') &&
		   ((*Part)[1] == L'C' || (*Part)[1] == L'c') &&
		    (*Part)[2] == L'=')
			break;
	while(*++Part);

	if(!*Part) //too far
		RaiseTextError(L"DN contains no domain components (e.g. DC=com)");

	auto_buffer Domain;
	BOOLEAN FirstElement = TRUE;

	do
	{
		if(!FirstElement) //only need dots after first element
			Domain.append(L".", sizeof(WCHAR));
		else
			FirstElement = FALSE;

		Domain.append(*Part + 3, wcslen(*Part + 3) * sizeof(WCHAR)); //3 being the third character: after DC=
	}
	while(*++Part);

	Domain.append(L"\0", sizeof(WCHAR));

	return Domain.release_as<WCHAR>();
}

void SetLdapOpts(LDAP *ld, SDE_OPTS *Opts)
{
	for(ULONG i = OiLdapOpts; i < Opts->OptCount; i++)
	{
		SDE_OPT *Opt = &Opts->Opts[i];
		UCHAR OptNum = static_cast<UCHAR>((Opt->Flags >> LDAP_OPTNUM_SHIFT) & 0xFF);

		if(Opt->Flags & SDE_OPT_SET)
			switch(Opt->Type)
			{
			case SptString:
				ldap_set_option(ld, OptNum, Opt->String);
				break;

			case SptInt32:
			case SptInt32Hex:
				ldap_set_option(ld, OptNum, &Opt->Int32);
				break;

			case SptEnum:
			case SptMultiEnumOr:
				ldap_set_option(ld, OptNum, &Opt->Enum->Choice);
				break;
			}
	}
}

class LdapConnection : public auto_refctd_std_base
{
public:
	auto_refct<NtdsLdapModuleBase> Module;
	SDE_OPTS *Opts;
	auto_free<LDAP, &free_ldap_connection> ld;
	auto_szpod<WCHAR> DomainDc;
	auto_szpod<WCHAR> ForestDc;
	ULONG TimeLimit;
	ULONG SizeLimit;

	//Win32ServerErrorLock is used to avoid the race concerning LDAP_OPT_SERVER_EXT_ERROR
	//held shared during operations that DO NOT use LDAP_OPT_SERVER_EXT_ERROR but will overwrite the value
	//held exclusively during operations that DO
	RTL_RESOURCE Win32ServerErrorLock; 

private:
	void Connect(PCWSTR Host)
	{
		OpText Ot(L"Connecting to ldap server", Host);
		SDE_OPT *Opt = Opts->Opts;

		ld = auto_take_(RaiseOnLdapErrorT<LDAP *>(ldap_init(const_cast<PWCHAR const>(Host), LDAP_PORT), NULL));

		SetLdapOpts(ld, Opts);

		RaiseOnLdapError(ldap_bind_s(ld, const_cast<PWCHAR const>(Opt[OiWho].String), NULL, c_cast<ULONG>(Opt[OiAuth].Enum->Choice)));

		Ot;
	}

	void Init()
	{
		TimeLimit = 2;
		SizeLimit = 0;
		RtlInitializeResource(&Win32ServerErrorLock);
	}

public:
	LdapConnection(NtdsLdapModuleBase *Module, SDE_GENERIC_OPTS const &GenOpts, UCHAR Privileges, PCWSTR ForDn = NULL OPTIONAL, PCWSTR ExplicitDomain = NULL OPTIONAL) :
		Module(Module),
		Opts(&Module->ModOpts) 
	{
		Init();

		SDE_OPT *Opt = Opts->Opts;

		PCWSTR Host = Opt[OiHost].Flags & SDE_OPT_SET ? Opt[OiHost].String : GenOpts.Host;
		auto_ptr<DOMAIN_CONTROLLER_INFO, Netapi::NetapiHeap> DcInfo;
		auto_szpod<WCHAR> Domain = ExplicitDomain ? ExplicitDomain : (Opt[OiDomain].Flags & SDE_OPT_SET ? Opt[OiDomain].String : GenOpts.Domain);
		PCWSTR Forest = Opt[OiForest].Flags & SDE_OPT_SET ? Opt[OiForest].String : NULL;
		
		if(!Domain && ForDn)
			Domain = DnToDomainDns(ForDn); //will not return null

		if(Opt[OiDomainDc].Flags & SDE_OPT_SET)
			DomainDc = Opt[OiDomainDc].String;
		else
			DomainDc = DomainDnsToDc(Domain);

		if(Opt[OiForestDc].Flags & SDE_OPT_SET)
			ForestDc = Opt[OiForestDc].String;
		else
			ForestDc = DomainDnsToDc(Forest); //if null now, lookup performed below

		if(!ForestDc || !Host)
		{
			OpText Ot(L"Trying to find a LDAP domain controller", Domain);

			RaiseOnWin32Error(DsGetDcName(GenOpts.Host, Domain, NULL, NULL, static_cast<ULONG>(Opt[OiDsnameopts].Enum->Choice), DcInfo.assign_indirect()));

			if(!ForestDc)
				ForestDc = DomainDnsToDc(DcInfo->DnsForestName);

			if(!Host)
				Host = Netapi::SkipAnyLeadingBackslashes(DcInfo->DomainControllerName);

			Ot;
		}
		
		Connect(Host);
	}

	LdapConnection(NtdsLdapModuleBase *Module, PCWSTR Host, PCWSTR DomainName, PCWSTR ForestName) :
		Module(Module),
		Opts(&Module->ModOpts)
	{
		Init();

		DomainDc = DomainDnsToDc(DomainName);
		ForestDc = DomainDnsToDc(ForestName);

		Connect(Host);
	}

	//returns NULL on failure
	static LdapConnection *TryOpenConnectionCrit(NtdsLdapModuleBase *Module, SDE_GENERIC_OPTS const &GenOpts)
	{
		SDE_OPTS *LdapOpts = &Module->ModOpts;

		//is the caller sure that this is a NTDS domain?
		//if they specified a host just for ldap, they must be
		if(LdapOpts->Opts[OiHost].Flags & SDE_OPT_SPECIFIED)
			return new LdapConnection(Module, GenOpts, 0);
		else //let's find out if there really is a NTDS domain
		{
			OpText Ot(L"Determining if host is an NTDS domain controller");

			auto_free<DSROLE_PRIMARY_DOMAIN_INFO_BASIC, &Netapi::free_DsRoleBasic> DsDomainInfo;
			RaiseOnWin32Error(DsRoleGetPrimaryDomainInformation(GenOpts.Host, DsRolePrimaryDomainInfoBasic, c_cast<PBYTE *>(DsDomainInfo.assign_indirect())));

			ULONG DiFlags = Netapi::GetPrimaryDomainFlagsFromDsRoleInfo(DsDomainInfo);

			if(DiFlags & Lookup::DOMAIN_INFO_LOCALDOMAIN)
				return NULL;
			else
				return new LdapConnection(Module, GenOpts, 0, NULL, DsDomainInfo->DomainNameDns ? DsDomainInfo->DomainNameDns : DsDomainInfo->DomainNameFlat);

			Ot;
		}
	}

	static LdapConnection *TryOpenConnection(NtdsLdapModuleBase *Module, SDE_GENERIC_OPTS const &GenOpts)
	{
		__try
		{
			return TryOpenConnectionCrit(Module, GenOpts);
		}
		__except(exception_code() == STATUS_VCPP_MOD_NOT_FOUND ||
			     exception_code() == STATUS_VCPP_PROC_NOT_FOUND
				 ? EXCEPTION_EXECUTE_HANDLER : EXCEPTION_CONTINUE_SEARCH)
		{//apparently, the proper NTDS support libraries or functions are not available. The proper return is null, i.e. functionally impossible
			return NULL;
		}
	}

	~LdapConnection()
	{
		RtlDeleteResource(&Win32ServerErrorLock);
	}
};

//holds a LdapConnection::Win32ServerErrorLock for the lifetime of this object
struct LcWin32ServerErrorLock
{
	LdapConnection *Lc;

	//set Exclusive to true to preserve value of LDAP_OPT_SERVER_EXT_ERROR
	LcWin32ServerErrorLock(LdapConnection *Lc, BOOLEAN Exclusive = FALSE) :
		Lc(Lc)
	{
		if(Exclusive)
			RtlAcquireResourceExclusive(&Lc->Win32ServerErrorLock, TRUE);
		else
			RtlAcquireResourceShared(&Lc->Win32ServerErrorLock, TRUE);
	}

	~LcWin32ServerErrorLock()
	{
		RtlReleaseResource(&Lc->Win32ServerErrorLock);
	}

	NO_ASSIGNMENT_OPERATOR(LcWin32ServerErrorLock)
	NO_COPY_CONSTRUCTOR(LcWin32ServerErrorLock)
};

enum SearchBaseType
{
	SbtLiteral, //use root as specified
	SbtDomain, //relative to domain
	SbtForest //relative to forest
};

ULONG __cdecl StartLdapSearchExV(LdapConnection *lc, ULONG Scope, PWSTR Attrs[], LDAPControl **ServerControls OPTIONAL, LDAPControl **ClientControls OPTIONAL, PCWSTR Root, SearchBaseType RootType, PCWSTR FilterFormat, va_list args)
{
	auto_szpod<WCHAR> FilterText;
	WCHAR FullRoot[MAX_STRING_CHARS];

	if(Root && Root[0])
		switch(RootType)
		{
		case SbtLiteral:
			wcscpy_s(FullRoot, Root);
			break;

		case SbtDomain:
			RaiseOnStdCError(swprintf_s(FullRoot, MAX_STRING_CHARS, L"%s,%s", Root, lc->DomainDc), -1);
			break;

		case SbtForest:
			RaiseOnStdCError(swprintf_s(FullRoot, MAX_STRING_CHARS, L"%s,%s", Root, lc->DomainDc), -1);
			break;
		}
	else
		switch(RootType)
		{
		case SbtLiteral:
			FullRoot[0] = '\0';
			break;

		case SbtDomain:
			wcscpy_s(FullRoot, lc->DomainDc);
			break;

		case SbtForest:
			wcscpy_s(FullRoot, lc->ForestDc);
			break;
		}

	SIZE_T FilterBufferLen = MAX_STRING_CHARS;

	while(TRUE)
	{
		FilterText.realloc_ct(FilterBufferLen);
		
		if(_vswprintf_p(FilterText, FilterBufferLen, FilterFormat, args) != -1)
			break;
		
		FilterBufferLen *= 2;
	}

	ULONG Msg;

	{
		LcWin32ServerErrorLock(lc, FALSE);

		RaiseOnLdapError(ldap_search_ext(lc->ld, FullRoot, Scope, FilterText, Attrs, FALSE, ServerControls, ClientControls, lc->TimeLimit, lc->SizeLimit, &Msg)); //return is connection # or -1 for error
	}

	return Msg;
}

ULONG __cdecl StartLdapSearchEx(LdapConnection *lc, ULONG Scope, PWSTR Attrs[], LDAPControl **ServerControls OPTIONAL, LDAPControl **ClientControls OPTIONAL, PCWSTR Root, SearchBaseType RootType, PCWSTR FilterFormat, ...)
{
	va_list args;
	va_start(args, FilterFormat); 

	return StartLdapSearchExV(lc, Scope, Attrs, ServerControls, ClientControls, Root, RootType, FilterFormat, args);
}

ULONG __cdecl StartLdapSearch(LdapConnection *lc, ULONG Scope, PWSTR Attrs[], PCWSTR Root, SearchBaseType RootType, PCWSTR FilterFormat, ...)
{
	va_list args;
	va_start(args, FilterFormat); 

	return StartLdapSearchExV(lc, Scope, Attrs, NULL, NULL, Root, RootType, FilterFormat, args);
}

void AwaitSearchResult(LdapConnection *lc, ULONG Msg, auto_free<LDAPMessage, &free_ldap_message> &SearchResult OUT)
{
	LcWin32ServerErrorLock(lc, FALSE);
	RaiseOnLdapErrorT<ULONG>(ldap_result(lc->ld, Msg, LDAP_MSG_ALL, NULL, SearchResult.assign_indirect()), (ULONG)-1);
}

template <class T> auto_take<T> AwaitSingleSearchResult(LdapConnection *lc, ULONG Msg, PCWSTR Attr)
{
	auto_free<LDAPMessage, &free_ldap_message> SearchResult;
	LcWin32ServerErrorLock(lc, FALSE);

	RaiseOnLdapErrorT<ULONG>(ldap_result(lc->ld, Msg, LDAP_MSG_ONE, NULL, SearchResult.assign_indirect()), (ULONG)-1);

	LDAPMessage *FirstEntry = ldap_first_entry(lc->ld, SearchResult);

	if(FirstEntry)
	{
		auto_free<berval *, free_ldap_value_len> Values = auto_take_(ldap_get_values_len(lc->ld, FirstEntry, const_cast<PWSTR>(Attr)));

		auto_szpod<T> Ret;

		if(Values)
			Ret = reinterpret_cast<T *>(Values[0]->bv_val);

		return Ret.release();
	}

	return auto_take<T>(NULL);
}

template <> auto_take<WCHAR> AwaitSingleSearchResult<WCHAR>(LdapConnection *lc, ULONG Msg, PCWSTR Attr)
{
	auto_free<LDAPMessage, &free_ldap_message> SearchResult;
	LcWin32ServerErrorLock(lc, FALSE);

	RaiseOnLdapErrorT<ULONG>(ldap_result(lc->ld, Msg, LDAP_MSG_ONE, NULL, SearchResult.assign_indirect()), (ULONG)-1);
	
	LDAPMessage *FirstEntry = ldap_first_entry(lc->ld, SearchResult);

	if(FirstEntry)
	{
		auto_free<PWSTR, free_ldap_value> Values = auto_take<PWSTR>(ldap_get_values(lc->ld, FirstEntry, const_cast<PWSTR>(Attr)));

		auto_szpod<WCHAR> Ret;

		if(Values)
			Ret = Values[0];

		return Ret.release();
	}

	return auto_take<WCHAR>(NULL);
}

}//namespace Ldap